"""
eDisk data reduction script
This script was written for CASA 6.1.1/6.2
Originally derived from DSHARP reduction scripts

Datasets calibrated (in order of date observed):
SB1: 

LB1: 

reducer: 
"""

""" Starting matter """
sys.path.append('/lustre/cv/users/jtobin/analysis_scripts/')
import analysisUtils as au
import analysisUtils as aU
import string
import os
import glob
import numpy as np
import sys
import pickle
execfile('reduction_utils3.py', globals())


###############################################################
################ SETUP/METADATA SPECIFICATION #################
################ USERS NEED TO SET STUFF HERE #################
###############################################################

### Use MPI CASA for faster imaging (start casa with mpicasa -n XX CASA; where XX is the number of processes >= 2)
parallel=True

### if True, can run script non-interactively if later parameters properly set
skip_plots = False	

### Add field names (corresponding to the field in the MS) here and prefix for 
### filenameing (can be different but try to keep same)
### Only make different if, for example, the field name has a space
field   = 'V883_Ori'
prefix  = 'V883_Ori' 

### always include trailing slashes!!
WD_path = '/lustre/cv/users/jtobin/V883Ori_water/HDO/'
SB_path = WD_path+'SB/'
LB_path = WD_path+'LB/'

### scales for multi-scale clean
SB_scales = [0, 5] #[0, 5, 10, 20]
LB_scales = [0, 5, 30]  #[0, 5, 30, 100, 200]

### automasking parameters for very extended emission
#sidelobethreshold=2.0
#noisethreshold=3.75
#lownoisethreshold=1.0
#smoothfactor=2.0
### automasking parameters for compact emission (uncomment to use)
sidelobethreshold=2.0
noisethreshold=4.0
lownoisethreshold=1.0 
smoothfactor=2.0

#read in final data_params from continuum to ensure we get the phase centers for each MS
with open(prefix+'.pickle', 'rb') as handle:
    data_params = pickle.load(handle)


###############################################################
#################### SHIFT PHASE CENTERS ######################
###############################################################
selectedVis='vis'
#selectedVis='vis_shift'

if selectedVis == 'vis_shift':
   for i in data_params.keys():
      data_params[i]['vis_shift']=prefix+'_'+i+'_shift.ms'
      os.system('rm -rf '+data_params[i]['vis_shift']+'*')
      fixvis(vis=data_params[i]['vis'], outputvis=data_params[i]['vis_shift'], 
         field=data_params[i]['field'], 
         phasecenter='J2000 '+data_params[i]['phasecenter'])
      fixplanets(vis=data_params[i]['vis_shift'], field=data_params[i]['field'], 
         direction=data_params[i]['common_dir'])

###############################################################
############### SCALE DATA RELATIVE TO ONE EB #################
###############################################################

### Uses scaling from continuum data
for i in data_params.keys():
   rescale_flux(data_params[i][selectedVis].replace(WD_path,''), [data_params[i]['gencal_scale']])
   data_params[i]['vis_rescaled']=data_params[i][selectedVis].replace('.ms','_rescaled.ms')

with open(prefix+'.pickle', 'wb') as handle:
    pickle.dump(data_params, handle, protocol=pickle.HIGHEST_PROTOCOL)

###############################################################
##### APPLY SELF-CALIBRATION SOLUTIONS TO LINE DATA ###########
###############################################################

### Gain tables and spw mapping saved to data dictionaries during selfcal and used as arguments here
for i in data_params.keys():
   n_tables=len(data_params[i]['selfcal_tables'])
   interp_list=['linearPD']*n_tables
   applycal(vis=data_params[i]['vis_rescaled'], spw='', 
         gaintable=data_params[i]['selfcal_tables'],spwmap=data_params[i]['selfcal_spwmap'], interp=interp_list, 
         calwt=True, applymode='calonly')
   split(vis=data_params[i]['vis_rescaled'],outputvis=data_params[i]['vis_rescaled'].replace('.ms','.ms.selfcal'),datacolumn='corrected')
   data_params[i]['vis_selfcal']=data_params[i]['vis_rescaled'].replace('.ms','.ms.selfcal')
   ### cleanup
   os.system('rm -rf '+data_params[i]['vis_rescaled'])
   if selectedVis=='vis_shift':
      os.system('rm -rf '+data_params[i]['vis_shift'])

with open(prefix+'.pickle', 'wb') as handle:
    pickle.dump(data_params, handle, protocol=pickle.HIGHEST_PROTOCOL)


###############################################################
################## DO CONTINUUM SUBTRACTION ###################
###############################################################

### Get channels to exclude for continuum fitting (same as the ones 
### we flagged for doing making continuum MS)
for i in data_params.keys():
   flagchannels_string = get_flagchannels(data_params[i], prefix)
   print(flagchannels_string)

   ### Get spws for argument list to uvcontsub
   spws_string = get_contsub_spws_indivdual_ms(data_params[i], prefix,only_cont_spws=True)
   print(spws_string)

   ### Run uvcontsub on combined, self-cal applied dataset; THIS WILL TAKE MANY HOURS PER EB
   contsub(data_params[i]['vis_selfcal'], prefix, spw=spws_string,flagchannels=flagchannels_string,excludechans=True)
   os.system('rm -rf '+prefix+'_'+i+'_spectral_line.ms')  ### remove existing spectral line MS if present
   os.system('mv '+data_params[i]['vis_selfcal'].replace('.selfcal','.selfcal.contsub')+' '+prefix+'_'+i+'_spectral_line.ms')
   os.system('rm -rf '+data_params[i]['vis_selfcal'])
   data_params[i]['vis_contsub']=prefix+'_'+i+'_spectral_line.ms'

with open(prefix+'.pickle', 'wb') as handle:
    pickle.dump(data_params, handle, protocol=pickle.HIGHEST_PROTOCOL)

###############################################################
################ TAR UP FINAL CONTSUBBED DATA #################
###############################################################

for i in data_params.keys():
   if 'SB' in i:
      os.system('rm -rf '+data_params[i]['vis_contsub']+'.tgz')
      os.system('tar czf '+data_params[i]['vis_contsub']+'.tgz '+data_params[i]['vis_contsub'])

###############################################################
############ RUN A FINAL SPECTRAL LINE IMAGE SET ##############
###############################################################


### generate list of MS files to image
vislist=[]
for i in data_params.keys():
   if 'SB' in i:
      vislist.append(data_params[i]['vis_contsub'])

### Dictionary defining the spectral line imaging parameters.

image_list = {
        "HDO-225.535GHz":dict(chanstart='-11km/s', chanwidth='0.162km/s',
            nchan=185, linefreq='225.89672000GHz', linespw='3',
            robust=[0.5,2.0]),
        "HDO-241.558GHz":dict(chanstart='-11km/s', chanwidth='0.151km/s',
            nchan=200, linefreq='241.56155000GHz', linespw='8', 
            robust=[0.5,2.0]),
        "HDO-225.535GHz-0.3kms":dict(chanstart='-11km/s', chanwidth='0.3km/s',
            nchan=100, linefreq='225.896720000GHz', linespw='3',
            robust=[0.5,2.0]),
        "HDO-241.558GHz-0.3kms":dict(chanstart='-11km/s', chanwidth='0.3km/s',
            nchan=100, linefreq='241.56155000GHz', linespw='8', 
            robust=[0.5,2.0]),
        "C17O-224.711GHz":dict(chanstart='-11km/s', chanwidth='0.163km/s',
            nchan=185, linefreq='224.71474380GHz', linespw='2',
            robust=[0.5,2.0]),
        "H2CO-225.694GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='225.69777250GHz', linespw='0',
            robust=[2.0]),
        "CH3OD-226.535GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='226.53867400GHz', linespw='1',
            robust=[2.0]),
        "C17O-224.711GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='224.71474380GHz', linespw='2',
            robust=[2.0]),
        "HDO-225.535GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='225.896720000GHz', linespw='3',
            robust=[2.0]),
        "CN-225.656GHz-fullcube":dict(chanstart='', chanwidth='244.14kHz',
            nchan=-1, linefreq='226.65955840GHz', linespw='4',
            robust=[2.0]),
        "CN-225.871GHz-fullcube":dict(chanstart='', chanwidth='244.14kHz',
            nchan=-1, linefreq='226.87589600GHz', linespw='5',
            robust=[2.0]),
        "wide-2GHz-240.496GHz-fullcube":dict(chanstart='', chanwidth='1953.124kHz',
            nchan=-1, linefreq='240.496GHz', linespw='6',
            robust=[2.0]),
        "CH3OH-241.846GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='241.84360400GHz', linespw='7',
            robust=[2.0]),
        "HDO-241.558GHz-fullcube":dict(chanstart='', chanwidth='122.07kHz',
            nchan=-1, linefreq='241.56155000GHz', linespw='8', 
            robust=[2.0]),
 }

### Loop through the spectral line images and make images.

cellsize='0.0075arcsec'
imsize=768
for line in image_list:
    for robust in image_list[line]["robust"]:
        imagename = prefix+f'_SB_'+line+'_robust_'+str(robust)

        sigma = get_sensitivity(data_params, specmode='cube', \
                spw=[image_list[line]["linespw"]], chan=450,robust=robust,imsize=imsize,cellsize=cellsize)

        tclean_spectral_line_wrapper(vislist, imagename,
                image_list[line]["chanstart"], image_list[line]["chanwidth"], 
                image_list[line]["nchan"], image_list[line]["linefreq"], 
                image_list[line]["linespw"], SB_scales, threshold=3.0*sigma,
                imsize=imsize, cellsize=cellsize,robust=robust, 
                sidelobethreshold=sidelobethreshold, noisethreshold=noisethreshold,
                lownoisethreshold=lownoisethreshold,smoothfactor=smoothfactor,parallel=parallel,
                phasecenter=data_params['SB1']['common_dir'].replace('J2000','ICRS'))



###############################################################
################ CLEANUP AND FITS CONVERSION ##################
###############################################################


import glob
### Remove extra image products
os.system('rm -rf *.residual* *.psf* *.model* *dirty* *.sumwt* *.gridwt* *.workdirectory')

### Remove fits files and pbcor files from previous iterations. 
os.system("rm -rf *.pbcor* *.fits") 

imagelist=glob.glob('*.image') + glob.glob('*.image.tt0')
for image in imagelist:
   impbcor(imagename=image,pbimage=image.replace('image','pb'),outfile=image.replace('image','pbcor'))
   exportfits(imagename=image.replace('image','pbcor'),fitsimage=image.replace('image','pbcor')+'.fits',overwrite=True,dropdeg=True)
   exportfits(imagename=image,fitsimage=image+'.fits',overwrite=True,dropdeg=True)

imagelist=glob.glob('*.mask')
for image in imagelist:
   exportfits(imagename=image,fitsimage=image+'.fits',overwrite=True,dropdeg=True)
   os.system('gzip '+image+'.fits')

os.system('rm -rf *initcont*.pb')
imagelist=glob.glob('*.pb') + glob.glob('*.pb.tt0')
for image in imagelist:
   exportfits(imagename=image,fitsimage=image+'.fits',overwrite=True,dropdeg=True)
   os.system('gzip '+image+'.fits')


### Remove rescaled selfcal MSfiles
os.system('rm -rf *rescaled.ms.*')
os.system('rm -rf scale*')

### Make a directory to put the final products
os.system('rm -rf export')
os.system('mkdir export')
os.system('mv *.fits export/')
os.system('mv *.tgz export/')
os.system('mv *.pdf export/')

